package org.korosoft.rusalad.syncnrun;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.korosoft.rusalad.syncnrun.command.Answer;
import org.korosoft.rusalad.syncnrun.command.Command;
import org.korosoft.rusalad.syncnrun.command.FailAnswer;
import org.korosoft.rusalad.syncnrun.command.MessageDrivenCommand;
import org.korosoft.rusalad.syncnrun.messaging.Message;
import org.korosoft.rusalad.syncnrun.messaging.MessageProcessor;
import org.korosoft.rusalad.syncnrun.messaging.MessageRoute;
import org.korosoft.rusalad.syncnrun.messaging.Messenger;
import org.korosoft.rusalad.syncnrun.tools.CountUpAndDownLatch;
import org.korosoft.rusalad.syncnrun.tools.SocketFactory;
import org.korosoft.rusalad.syncnrun.tools.Utils;

import java.io.*;
import java.net.Socket;
import java.net.SocketException;
import java.util.Collection;
import java.util.HashMap;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * Syncnrun RPC client.
 *
 * @author Dmitry Korotkov
 * @since 1.0
 */
public class Client {
    private static final Log log = LogFactory.getLog(Client.class);

    private static final CommandWithId shutdownCommand = new CommandWithId(null, null, 0);
    private static final String paddingObject = Utils.getPaddingObject(256);
    private static final long paddingIntervalMs = 50;

    public static final int WRITE_QUEUE_CAPACITY = 128;
    private final SocketFactory socketFactory;
    private final String host;
    private final int port;
    private final int defaultAttemptsCount;

    private final ConcurrentMap<UUID, CountUpAndDownLatch> latches = new ConcurrentHashMap<UUID, CountUpAndDownLatch>();
    private final ConcurrentMap<UUID, Answer> answers = new ConcurrentHashMap<UUID, Answer>();
    private final BlockingQueue<CommandWithId> commands = new ArrayBlockingQueue<CommandWithId>(WRITE_QUEUE_CAPACITY);
    private final Semaphore reactorSemaphore = new Semaphore(1);

    public Client(SocketFactory socketFactory, String host, int port) throws IOException {
        this.socketFactory = socketFactory;
        this.host = host;
        this.port = port;
        this.defaultAttemptsCount = 5;
    }

    public Client(SocketFactory socketFactory, String host, int port, int defaultAttemptsCount) throws IOException {
        this.socketFactory = socketFactory;
        this.host = host;
        this.port = port;
        this.defaultAttemptsCount = defaultAttemptsCount;
    }

    public <A extends Answer> A execCommandSynchronously(Command<A> command) {
        return execCommandSynchronously(command, defaultAttemptsCount);
    }

    public <A extends Answer> A execCommandSynchronously(Command<A> command, int attempts) {
        UUID uuid = postCommand(command, attempts);
        try {
            latches.get(uuid).await();
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        }
        latches.remove(uuid);
        Answer answer = answers.remove(uuid);
        if (answer instanceof FailAnswer) {
            throw new RuntimeException(((FailAnswer) answer).exception);
        }
        //noinspection unchecked
        return (A) answer;
    }

    public void stop() throws InterruptedException {
        if (commands.size() > 0) {
            log.debug("Waiting for all commands to execute...");
        }
        while (commands.size() > 0) {
            Thread.sleep(1000);
        }
        log.debug("Command queue is empty. Stopping client reactor.");
        commands.put(shutdownCommand);
        reactorSemaphore.acquire();
        log.debug("Client reactor stopped.");
    }

    private UUID postCommand(Command command, int attempts) {
        UUID uuid = UUID.randomUUID();
        latches.put(uuid, new CountUpAndDownLatch(1));
        CommandWithId cid = new CommandWithId(uuid, command, attempts);
        try {
            commands.put(cid);
        } catch (InterruptedException e) {
            throw new RuntimeException();
        }
        checkReactorIsRunning();
        return uuid;
    }

    private void checkReactorIsRunning() {
        if (reactorSemaphore.tryAcquire()) {
            new Reactor().start();
        }
    }

    private class Reactor {

        private final Map<UUID, CommandWithId> sentCommands = new HashMap<UUID, CommandWithId>();
        private Socket socket;
        private Thread reader;
        private StreamWriter writer;
        private Messenger messenger;

        private void start() {
            log.debug("Starting client reactor");
            new Thread(new Runnable() {
                @Override
                public void run() {
                    try {
                        for (; ; ) {
                            CommandWithId command = commands.take();
                            if (command == shutdownCommand) {
                                log.debug("Client reactor is shutting down...");
                                synchronized (Reactor.this) {
                                    socket.close();
                                    writer.close();
                                    reader.join();
                                }
                                break;
                            }
                            if (log.isTraceEnabled()) {
                                log.trace(String.format("Processing command %s", command.toString()));
                            }
                            try {
                                processCommand(command);
                            } catch (Exception e) {
                                commands.put(command);
                                int attempts = command.attemptsLeft.decrementAndGet();
                                if (attempts <= 0) {
                                    log.error(String.format("No attempts left for command %s. Exiting with failure.", command.command));
                                    System.exit(1);
                                }
                                log.warn(String.format("Failed to process command. Command will be resent. Attempts left: %d. Error: %s", attempts, e.getMessage()));
                                if (log.isTraceEnabled()) {
                                    log.trace(e.getMessage(), e);
                                }
                            }
                        }
                    } catch (Exception e) {
                        throw new RuntimeException(e);
                    } finally {
                        reactorSemaphore.release();
                    }
                }
            }, "client-command-sender").start();
        }

        private void processCommand(CommandWithId command) throws InterruptedException, IOException {
            synchronized (this) {
                if (socket == null || !socket.isConnected() || socket.isClosed()) {
                    reconnect(true);
                }

                if (log.isDebugEnabled()) {
                    log.debug(String.format("Sending command ID=%s: '%s'", command.uuid, command.command));
                } else {
                    log.info(String.format("Sending command: '%s'", command.command));
                }
                if (command.command instanceof MessageDrivenCommand) {
                    final MessageDrivenCommand c = (MessageDrivenCommand) command.command;
                    for (UUID routeName : (Collection<UUID>) c.getRouteNames()) {
                        messenger.getRoute(routeName).setHotReceiver(new MessageDrivenCommandMsgProcessor(command.uuid, c));
                    }
                }
                writer.writeObject(command.uuid);
                writer.writeObject(command.command);
                sentCommands.put(command.uuid, command);
            }
        }

        private Semaphore reconnectSemaphore = new Semaphore(1);

        private void reconnect(boolean force) throws InterruptedException, IOException {
            if (!force) {
                if (!reconnectSemaphore.tryAcquire()) {
                    return;
                }
            } else {
                reconnectSemaphore.acquire();
            }
            try {
                synchronized (this) {
                    for (CommandWithId commandWithId : sentCommands.values()) {
                        int attempts = commandWithId.attemptsLeft.decrementAndGet();
                        if (attempts <= 0) {
                            log.error(String.format("No attempts left for command %s. Exiting with failure.", commandWithId.command));
                            System.exit(1);
                        }
                        log.debug(String.format("Command '%s' will be re-sent", commandWithId));
                        commands.put(commandWithId);
                    }
                    sentCommands.clear();
                    if (reader != null) {
                        reader.interrupt();
                        reader.join();
                    }
                    if (commands.size() > 0 || force) {
                        log.info(String.format("Connecting to server %s:%d", host, port));
                        if (this.socket != null) {
                            try {
                                socket.close();
                            } catch (IOException ignored) {
                            }
                        }
                        final Socket socket = socketFactory.createClientSocket(host, port);
                        log.info("Connected");
                        this.socket = socket;

                        reader = new Thread(new ClientReader(socket), "client-reader");
                        reader.start();

                        writer = new StreamWriter(socket.getOutputStream());
                        new Thread(writer, "client-writer").start();
                        if (messenger != null) {
                            messenger.sendEndOfTransfer();
                        }
                        messenger = new Messenger(new MessageProcessor() {
                            @Override
                            public void processMessage(MessageRoute route, Message message) {
                                synchronized (this) {
                                    if (log.isTraceEnabled()) {
                                        log.trace(String.format("Queue %s: sending message: %s", route.getName(), message));
                                    }
                                    writer.writeObject(route.getName());
                                    writer.writeObject(message);
                                }
                            }
                        });
                    }
                }
            } finally {
                reconnectSemaphore.release();
            }
        }


        private class ClientReader implements Runnable {
            private final Socket socket;
            private final ExecutorService executorService = Executors.newCachedThreadPool(new ThreadFactory() {
                private final AtomicInteger no = new AtomicInteger(0);

                @Override
                public Thread newThread(Runnable r) {
                    String name = String.format("client-command-handler-%d", no.incrementAndGet());
                    if (log.isTraceEnabled()) {
                        log.trace(String.format("Spawning thread %s", name));
                    }
                    return new Thread(r, name);
                }
            });

            private ClientReader(Socket socket) {
                this.socket = socket;
            }

            @Override
            public void run() {
                ObjectInputStream stream = null;
                try {
                    if (log.isTraceEnabled()) {
                        log.trace("Starting client reader");
                    }
                    stream = new ObjectInputStream(socket.getInputStream());
                    for (Object o = stream.readObject(); o != null; o = stream.readObject()) {
                        if (o instanceof String) {
                            continue;
                        }
                        if (log.isTraceEnabled()) {
                            log.trace(String.format("Received id: %s", o));
                        }
                        if (!(o instanceof UUID)) {
                            throw new RuntimeException("Expected UUID but received " + o.getClass().getName());
                        }
                        final UUID uuid = (UUID) o;
                        do {
                            o = stream.readObject();
                        } while (o instanceof String);
                        if (o instanceof Message) {
                            final Message message = (Message) o;
                            final Messenger messenger;
                            synchronized (Reactor.this) {
                                messenger = Reactor.this.messenger;
                            }
                            // The below code updated command latch so it can be finished only when all the messages are processed.
                            MessageProcessor hotReceiver = messenger.getRoute(uuid).getHotReceiver();
                            final MessageDrivenCommandMsgProcessor proc;
                            if (hotReceiver instanceof MessageDrivenCommandMsgProcessor) {
                                proc = ((MessageDrivenCommandMsgProcessor) hotReceiver);
                            } else {
                                proc = null;
                            }
                            if (proc != null) {
                                latches.get(proc.commandId).countUp();
                            }
                            executorService.execute(new Runnable() {
                                @Override
                                public void run() {
                                    try {
                                        messenger.getRoute(uuid).processIncoming(message);
                                    } catch (InterruptedException e) {
                                        throw new RuntimeException(e);
                                    } finally {
                                        if (proc != null) {
                                            latches.get(proc.commandId).countDown();
                                        }
                                    }
                                }
                            });
                            if (log.isTraceEnabled()) {
                                log.trace(String.format("Queue %s: received message: %s", uuid, message));
                            }
                            continue;
                        }
                        if (!(o instanceof Answer)) {
                            throw new RuntimeException("Expected Answer but received " + o.getClass().getName());
                        }
                        CommandWithId cmd = sentCommands.remove(uuid);
                        if (o instanceof FailAnswer) {
                            if (log.isDebugEnabled()) {
                                log.debug(String.format("For command id %s a failure answer was received. Server stack trace is provided.", uuid), ((FailAnswer) o).exception);
                            } else {
                                log.info(String.format("Received an answer for command '%s': '%s'", cmd.command, ((FailAnswer) o).exception.getMessage()));
                            }
                        } else {
                            if (log.isDebugEnabled()) {
                                log.debug(String.format("For command id %s an answer was received: %s", uuid, o));
                            } else {
                                log.info(String.format("Received an answer for command '%s': '%s'", cmd.command, o));
                            }
                        }
                        if (latches.containsKey(uuid)) {
                            answers.put(uuid, (Answer) o);
                            latches.get(uuid).countDown();
                        } else {
                            throw new RuntimeException("Unknown command ID received with answer " + o.getClass().getName());
                        }
                    }
                } catch (EOFException e) {
                    try {
                        log.warn(String.format("Server unexpectedly broken the connection. Reconnecting..."));
                        socket.close();
                        synchronized (this) {
                            reconnect(false);
                        }
                    } catch (Exception ignored) {
                    }
                } catch (SocketException e) {
                    if (socket.isClosed()) {
                        log.debug("Socket is closed. Client reader exited");
                        return;
                    } else {
                        try {
                            log.warn(String.format("Server connection broken."));
                            socket.close();
                            new Thread(new Runnable() {
                                @Override
                                public void run() {
                                    try {
                                        reconnect(false);
                                    } catch (Exception ignored) {
                                    }
                                }
                            }, "reconnecter").start();
                            return;
                        } catch (Exception ignored) {
                        }
                        if (log.isDebugEnabled()) {
                            log.error(e.getMessage(), e);
                        } else {
                            log.error(e.getMessage());
                        }
                        return;
                    }
                } catch (Exception e) {
                    try {
                        socket.close();
                        if (log.isDebugEnabled()) {
                            log.debug("Closing server connection.", e);

                        } else {
                            log.warn("Closing server connection.");
                        }

                        new Thread(new Runnable() {
                            @Override
                            public void run() {
                                try {
                                    reconnect(false);
                                } catch (Exception ignored) {
                                }
                            }
                        }, "reconnecter").start();
                        return;
                    } catch (Exception ignored) {
                    }
                    if (log.isDebugEnabled()) {
                        log.error(e.getMessage(), e);
                    } else {
                        log.error(e.getMessage());
                    }
                    return;
                } finally {
                    if (stream != null) {
                        try {
                            stream.close();
                        } catch (Exception ignored) {
                        }
                    }
                }
                if (messenger != null) {
                    messenger.sendEndOfTransfer();
                }
                try {
                    socket.close();
                } catch (Exception ignored) {
                }
            }
        }

        private class StreamWriter implements Runnable {
            private final BlockingQueue<Object> writeQueue = new ArrayBlockingQueue<Object>(WRITE_QUEUE_CAPACITY);
            private final ObjectOutputStream objectOutputStream;

            private StreamWriter(OutputStream outputStream) throws IOException {
                this.objectOutputStream = new ObjectOutputStream(outputStream);
            }

            public void writeObject(Object o) {
                try {
                    writeQueue.put(o);
                } catch (InterruptedException e) {
                    throw new RuntimeException();
                }
            }

            public void close() {
                try {
                    writeQueue.put(this);
                } catch (InterruptedException e) {
                    throw new RuntimeException();
                }
            }

            @Override
            public void run() {
                try {
                    boolean lastSentObjectIsPadding = false;
                    for (; ; ) {
                        final Object o;
                        if (lastSentObjectIsPadding) {
                            o = writeQueue.take();
                        } else {
                            o = writeQueue.poll(paddingIntervalMs, TimeUnit.MILLISECONDS);
                        }
                        if (o == null) {
                            if (log.isTraceEnabled()) {
                                log.trace("Sending padding object");
                            }
                            objectOutputStream.writeObject(paddingObject);
                            objectOutputStream.flush();
                            lastSentObjectIsPadding = true;
                            continue;
                        }
                        if (o == this) {
                            break;
                        }
                        if (log.isTraceEnabled()) {
                            log.trace(String.format("Sending object %s", o));
                        }
                        objectOutputStream.writeObject(o);
                        objectOutputStream.reset();
                        lastSentObjectIsPadding = false;
                    }
                    objectOutputStream.close();
                } catch (Exception e) {
                    close();
                    try {
                        reconnect(false);
                    } catch (Exception ignored) {
                    }
                    throw new RuntimeException(e);
                }
            }
        }

    }

    private static class CommandWithId {
        private final UUID uuid;
        private final Command command;
        private final AtomicInteger attemptsLeft;

        private CommandWithId(UUID uuid, Command command, int attempts) {
            this.uuid = uuid;
            this.command = command;
            this.attemptsLeft = new AtomicInteger(attempts);
        }

        @Override
        public String toString() {
            return String.format("[%s, %s]", uuid, command);
        }
    }

    private class MessageDrivenCommandMsgProcessor implements MessageProcessor {
        private final MessageDrivenCommand<?> command;
        private final UUID commandId;

        private MessageDrivenCommandMsgProcessor(UUID commandId, MessageDrivenCommand<?> command) {
            this.commandId = commandId;
            this.command = command;
        }

        @Override
        public void processMessage(MessageRoute route, Message message) {
            command.processMessageAtClient(route, message);
        }
    }
}
